{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "complete-runner",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.ensemble import GradientBoostingRegressor\n",
    "from sklearn.ensemble import ExtraTreesRegressor\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.experimental import enable_hist_gradient_boosting\n",
    "from sklearn.ensemble import HistGradientBoostingRegressor\n",
    "from sklearn.ensemble import AdaBoostRegressor\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "from sklearn.neural_network import MLPRegressor\n",
    "from sklearn.tree import DecisionTreeRegressor\n",
    "\n",
    "Cnt = 0;\n",
    "#为每个表编号\n",
    "Dict_table = {}\n",
    "Dict_table['t'] = 0\n",
    "Dict_table['mc'] = 1\n",
    "Dict_table['ci'] = 2\n",
    "Dict_table['mi'] = 3\n",
    "Dict_table['mi_idx'] = 4\n",
    "Dict_table['mk'] = 5\n",
    "\n",
    "# t.production_year\n",
    "# mi.info_type_id\n",
    "# t.kind_id\n",
    "# mi_idx.info_type_id\n",
    "# mk.keyword_id\n",
    "# ci.person_id\n",
    "# mc.company_id\n",
    "# mc.company_type_id\n",
    "# ci.role_id\n",
    "# 为训练集出现过的9个属性编号\n",
    "Dict_attr = {}\n",
    "Dict_attr['t.production_year'] = 0\n",
    "Dict_attr['mi.info_type_id'] = 1\n",
    "Dict_attr['t.kind_id'] = 2\n",
    "Dict_attr['mi_idx.info_type_id'] = 3\n",
    "Dict_attr['mk.keyword_id'] = 4\n",
    "Dict_attr['ci.person_id'] = 5\n",
    "Dict_attr['mc.company_id'] = 6\n",
    "Dict_attr['mc.company_type_id'] = 7\n",
    "Dict_attr['ci.role_id'] = 8\n",
    "\n",
    "Dict_comp = {}\n",
    "Dict_comp['>'] = 0\n",
    "Dict_comp['='] = 1\n",
    "Dict_comp['<'] = 2\n",
    "\n",
    "Dict_max = {}\n",
    "Dict_min = {}\n",
    "\n",
    "Self_add_x = []\n",
    "Self_add_y = []\n",
    "\n",
    "# load Max, Min information, 用来归一化\n",
    "with open('column_min_max_vals.csv', 'r') as csvfile:\n",
    "    Max_list = csv.reader(csvfile)\n",
    "    for Max_data in Max_list:\n",
    "        if (Max_data[0] == \"name\"):\n",
    "            continue\n",
    "        Dict_max[Max_data[0]] = Max_data[2] #每个属性的最大值和最小值信息存入字典\n",
    "        Dict_min[Max_data[0]] = Max_data[1]\n",
    "        \n",
    "        Cur_add = [] #额外增加几条语句，用上表中的cardinality信息\n",
    "\n",
    "        for i in range(42):\n",
    "            Cur_add.append(0.0)\n",
    "        Cur_add[Dict_table[Max_data[0].split('.')[0]]] = 1.0 #构造单表无查询的特征向量，结果即为单表数据数\n",
    "        Self_add_x.append(Cur_add)\n",
    "        Self_add_y.append(math.log(int(Max_data[3])) / math.log(459922927)) #459922927为train集中最大查询数量，这里先取log再归一化\n",
    "\n",
    "def Parse(Data, Is_train): #Is_train=1表示Data是训练集，否则为测试集\n",
    "    \n",
    "    Feature_vec = [] #42维的X变量\n",
    "    Feature_target = [] #Y值\n",
    "    \n",
    "    for Sql in Data:\n",
    "        Table_name = [] #Table的名字\n",
    "        Join_table = [] #连接部分\n",
    "        Condition = [] #条件\n",
    "        Target = 0 #Is_train=1时表示满足条件的数量\n",
    "        \n",
    "        Table_name = Sql[0].split(',')\n",
    "        Join_table = Sql[1].split(',')\n",
    "\n",
    "        Cur_cond = Sql[2].split(',')\n",
    "        Cur_Size = len(Cur_cond)\n",
    "        Cur_Cnt = 0\n",
    "        \n",
    "        while (Cur_Cnt < Cur_Size): #每三个为1组\n",
    "            Cur_list = []\n",
    "            Cur_list.append(Cur_cond[Cur_Cnt])\n",
    "            Cur_list.append(Cur_cond[Cur_Cnt+1])\n",
    "            Cur_list.append(Cur_cond[Cur_Cnt+2])\n",
    "            Condition.append(Cur_list)\n",
    "            Cur_Cnt = Cur_Cnt + 3\n",
    "        \n",
    "        if (Is_train == 1):\n",
    "            Target = int(Sql[3])\n",
    "\n",
    "        # Create vec\n",
    "        Cur_vec = [] #该条SQL语句的特征向量\n",
    "        for i in range(42): #不存在默认为补0\n",
    "            Cur_vec.append(0.0)\n",
    "        \n",
    "        #录入表信息，6维表向量，包含哪些表在哪个位置为1\n",
    "        pos = 0\n",
    "        for Cur_table_name in Table_name:\n",
    "            Cur_vec[Dict_table[Cur_table_name.split(' ')[-1]]] = 1.0\n",
    "        pos = pos + 6\n",
    "\n",
    "        #条件包含36维特征向量，每个属性占4个维度，且已在字典提前定义过属性的顺序\n",
    "        #（>,=,<,val）\n",
    "        #四维的前三维为独热向量，分别表示>,=,<,如果在该属性不存在条件则全为0\n",
    "        #四维的最后一维表示val\n",
    "        for i in Condition:\n",
    "            offset = Dict_attr[i[0]] * 4;\n",
    "            Cur_vec[pos + offset + Dict_comp[i[1]]] = 1.0\n",
    "            Cur_vec[pos + offset + 3] = (1.0 * int(i[2]) - 1.0* int(Dict_min[i[0]]) ) / (1.0 * int(Dict_max[i[0]]) - 1.0 * int(Dict_min[i[0]]))\n",
    "        #利用该属性最大最小值进行归一化处理\n",
    "        Feature_vec.append(Cur_vec)\n",
    "        if (Is_train == 1):\n",
    "            Feature_target.append(math.log(1.0*Target) / math.log(459922927)) #459922927为train集中最大查询数量，这里先取log再归一化\n",
    "            #取ln是为了让数据的变化更敏感，误差更小\n",
    "    return [Feature_vec, Feature_target]\n",
    "\n",
    "X_Data = []\n",
    "Y_Data = []\n",
    "\n",
    "#读取训练集数据\n",
    "with open('train.csv', 'r') as csvfile:\n",
    "    Data = csv.reader(csvfile, delimiter = '#')\n",
    "    X_Data = Parse(Data, 1)\n",
    "\n",
    "#读取需要预测的数据\n",
    "with open('test_without_label.csv', 'r') as csvfile:\n",
    "    Data = csv.reader(csvfile, delimiter = '#')\n",
    "    Y_Data = Parse(Data, 0)\n",
    "\n",
    "X_train = X_Data[0]\n",
    "Y_train = X_Data[1]\n",
    "\n",
    "#加上自己增加的几条单表边界条件\n",
    "CurLen = len(Self_add_x)\n",
    "for i in range(CurLen):\n",
    "    X_train.append(Self_add_x[i])\n",
    "    Y_train.append(Self_add_y[i])\n",
    "\n",
    "X_test = Y_Data[0]\n",
    "\n",
    "clf = HistGradientBoostingRegressor(max_iter = 500, max_depth=7) #建立模型\n",
    "clf.fit(X_train, Y_train) #对训练集进行机器学习\n",
    "\n",
    "Y_test = clf.predict(X_test) #对测试集进行测试\n",
    "\n",
    "Writeline = []\n",
    "Cur_Len = len(Y_test)\n",
    "for i in range(Cur_Len):\n",
    "    Writeline.append([i, int(math.exp(Y_test[i] * math.log(459922927)))]) #之前归一化，这里乘回来，再取e的exp还原真实预测数据\n",
    "\n",
    "with open(\"2019201412.csv\", \"w\", newline = '') as csvfile:\n",
    "    csv_writer = csv.writer(csvfile)\n",
    "    csv_writer.writerow(['Query ID', 'Predicted Cardinality']) #输出至本地文件\n",
    "    for i in Writeline:\n",
    "        csv_writer.writerow(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "macro-authority",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "referenced-announcement",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
